from openai import OpenAI
import json
import openai
import io
import pandas as pd

openai.api_key = "sk-6S0PtpNia71gjcfwSsDPsJ9mGqsVPr2XRQzAx1dHbJS7RW4t"
openai.api_base="https://chatapi.littlewheat.com/v1"


client = OpenAI(api_key=openai.api_key ,base_url=openai.api_base)

# Example dummy function hard coded to return the same weather
# In production, this could be your backend API or an external API
def sunwukong_function(data):
    """
    孙悟空算法函数，该函数定义了数据集计算过程
    :param data: 必要参数，表示带入计算的数据表，用字符串进行表示
    :return：sunwukong_function函数计算后的结果，返回结果为表示为JSON格式的Dataframe类型对象
    """
    data = io.StringIO(data)
    df_new = pd.read_csv(data, sep='\\s+', index_col=0)
    res = df_new['x1'] * 10
    return json.dumps(res.to_string())


df_str=pd.DataFrame({'x1':[1, 2], 'x2':[3, 4]}).to_string

def run_conversation():
    # Step 1: send the conversation and available functions to the model
    messages=[
    {"role": "system", "content": "数据集data：%s，数据集以字符串形式呈现" % df_str},
    {"role": "user", "content": "请在数据集data上执行孙悟空算法"}
        ]
    tools = [
    {
        "type": "function",
        "function": {"name": "sunwukong_function",
                      "description": "用于执行孙悟空算法函数，定义了一种特殊的数据集计算过程",
                      "parameters": {"type": "object",
                                     "properties": {"data": {"type": "string",
                                                             "description": "执行孙悟空算法的数据集"},
                                                   },
                                     "required": ["data"],
                                    },
                     }
    }]

    response = client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=messages,
        tools=tools,
        tool_choice="auto",  # auto is default, but we'll be explicit
    )
    response_message = response.choices[0].message
    tool_calls = response_message.tool_calls
    # Step 2: check if the model wanted to call a function
    if tool_calls:
        # Step 3: call the function
        # Note: the JSON response may not always be valid; be sure to handle errors
        available_functions = {
            "sunwukong_function": sunwukong_function,
        }  # only one function in this example, but you can have multiple
        messages.append(response_message)  # extend conversation with assistant's reply
        # Step 4: send the info for each function call and function response to the model
        for tool_call in tool_calls:
            function_name = tool_call.function.name
            function_to_call = available_functions[function_name]
            function_args = json.loads(tool_call.function.arguments)
            function_response = function_to_call(**function_args)
            messages.append(
                {
                    "tool_call_id": tool_call.id,
                    "role": "tool",
                    "name": function_name,
                    "content": function_response,
                }
            )  # extend conversation with function response
        second_response = client.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=messages,
        )  # get a new response from the model where it can see the function response
        return second_response
result=run_conversation()
print(result)